import cv2 
import shelve
import numpy as np

class easySample():
    """"
        通过文件路径--》各个模型的样本输入
    """
    sample=None
    def __init__(self) -> None:
        self.imgPath = r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/img"

    def get_sample(self,name,modelName):
        if modelName =="siamese_img":
            return self.get_siamese_img_sample(name)
        elif modelName =="functionSim":
            return self.get_functionSim_sample(name)
            
    def get_functionSim_sample(self,name):
        """
            路径输入默认是ida进行反汇编
            输出：
                adj 邻接矩阵 b*n*n 
                att 特征矩阵 b*n*d
                vtype 类型矩阵 b*n*3  3种类型
        """
        fileroot = name
        res={}
        res["adj"],res["att"],res["vtype"]=[],[],[]
        with shelve.open(fileroot) as file:
            cg=file["cg"]
            cgattr=file["cgattr"]
            try:
                functype=file["funcType"]
            except Exception as e:
                functype=file["functype"]
            try:
                funcs_id=file["func_id"]# name-->ind
            except Exception as e:
                funcs_id=file["funcs_id"]# name-->ind
        id_funcs={}
        for i in funcs_id.keys():
            id_funcs[funcs_id[i]]=i

        #   动态导入函数是[],需要处理一下,
        #   同时生成vtype
        tempData=[]
        for i in functype.keys():
            temp=[]
            if functype[i]=="local":
                temp=[1,0,0]
            elif functype[i]=='dynamic import':
                temp=[0,1,0]
                cgattr[i]=self.map_api_name_to_eight_bit_vector(id_funcs[i],8,20)
            else:
                temp=[0,0,1]
            tempData.append(temp)
        res["adj"]=np.array(cg)#这里是带边权的
        res["att"]=np.array(cgattr)#应该是8维，但是第一维字符串的提取似乎出现了问题，建议只是用后七维
        res["vtype"]=np.array(tempData)
        return res
        
    def get_siamese_img_sample(self,name):
        res=cv2.imread(self.imgPath+"/"+name+".png",cv2.IMREAD_UNCHANGED)
        return res
    
    def map_api_name_to_eight_bit_vector(self,function_name,embeddingSize,numLim):
        """
            将api转换成一个embeddingSize位的向量，每一位使用numLim取余
        """
        function_bytes = function_name.encode('utf-8')
        eight_bit_vector = [0] * embeddingSize
        for i, byte in enumerate(function_bytes):
            index = i % embeddingSize
            eight_bit_vector[index] += byte
            eight_bit_vector[index] %= numLim
        return eight_bit_vector

if __name__=="__main__":    
    a=easySample()
    b=a.get_sample("0a9a57bce84229c959a9d852e6cfaf14a163f16eaa6b47b131b52d3c32a2bf73_trojan","functionSim")
    print(b)